mod request_id;
mod server_time;

use argon2::{
    password_hash::{rand_core::OsRng, PasswordHasher, SaltString},
    Argon2, PasswordHash, PasswordVerifier,
};

use crate::models::response::{Response as ErrorResponse, ResponseFailed};
use crate::{errors::error::AppError, models::user::User};
use axum::{
    extract::{FromRequestParts, Query, Request, State},
    http::StatusCode,
    middleware::{from_fn, Next},
    response::{IntoResponse, Response},
    Router,
};
use axum_extra::{
    headers::{authorization::Bearer, Authorization},
    TypedHeader,
};
use core::fmt;
use jwt_simple::token;
use request_id::set_request_id;
use serde::Deserialize;
use server_time::ServerTimeLayer;
use tower::ServiceBuilder;
use tower_http::{
    compression::CompressionLayer,
    trace::{DefaultMakeSpan, DefaultOnRequest, DefaultOnResponse, TraceLayer},
    LatencyUnit,
};
use tracing::{warn, Level};

const REQUEST_ID_HEADER: &str = "x-request-id";
const SERVER_TIME_HEADER: &str = "x-server-time";
pub trait TokenVerify {
    type Error: fmt::Debug;
    fn verify(&self, token: &str) -> Result<User, Self::Error>;
}

#[derive(Debug, Deserialize)]
struct Params {
    token: String,
}

pub fn set_layer(app: Router) -> Router {
    app.layer(
        ServiceBuilder::new()
            .layer(
                TraceLayer::new_for_http()
                    .make_span_with(DefaultMakeSpan::new().include_headers(true))
                    .on_request(DefaultOnRequest::new().level(Level::INFO))
                    .on_response(
                        DefaultOnResponse::new()
                            .level(Level::INFO)
                            .latency_unit(LatencyUnit::Micros),
                    ),
            )
            .layer(CompressionLayer::new().gzip(true).br(true).deflate(true))
            .layer(from_fn(set_request_id))
            .layer(ServerTimeLayer),
    )
}

pub fn hash_password(password: &str) -> Result<String, AppError> {
    let salt = SaltString::generate(&mut OsRng);
    // Argon2 with default params (Argon2id v19)
    let argon2 = Argon2::default();
    // Hash password to PHC string ($argon2id$v=19$...)
    let password_hash = argon2
        .hash_password(password.as_bytes(), &salt)?
        .to_string();

    Ok(password_hash)
}

pub fn verify_password(password: &str, password_hash: &str) -> Result<bool, AppError> {
    let argon2 = Argon2::default();
    let password_hash = PasswordHash::new(password_hash)?;
    let is_valid = argon2
        .verify_password(password.as_bytes(), &password_hash)
        .is_ok();
    Ok(is_valid)
}

pub async fn verify_token<T>(State(state): State<T>, req: Request, next: Next) -> Response
where
    T: TokenVerify + Clone + Send + Sync + 'static,
{
    let (mut parts, body) = req.into_parts();
    let token = match TypedHeader::<Authorization<Bearer>>::from_request_parts(&mut parts, &state)
        .await
    {
        Ok(TypedHeader(Authorization(bearer))) => bearer.token().to_string(),
        Err(e) => {
            if e.is_missing() {
                match Query::<Params>::from_request_parts(&mut parts, &state).await {
                    Ok(params) => params.token.clone(),
                    Err(e) => {
                        let msg = format!("parse query params failed: {}", e);
                        warn!(msg);
                        return ErrorResponse::<()>::failed(600, "Bad Request".to_string())
                            .into_response();
                        // return (StatusCode::UNAUTHORIZED, msg).into_response();
                    }
                }
            } else {
                let msg = format!("parse Authorization header failed: {}", e);
                warn!(msg);
                return ErrorResponse::<()>::failed(600, "Bad Request".to_string()).into_response();
            }
        }
    };

    let req = match state.verify(&token) {
        Ok(user) => {
            let mut req = Request::from_parts(parts, body);
            req.extensions_mut().insert(user);
            req
        }
        Err(e) => {
            let msg = format!("verify token failed: {:?}", e);
            warn!(msg);
            return ErrorResponse::<()>::failed(600, "Bad Request".to_string()).into_response();
        }
    };

    next.run(req).await
}
